% this file performs Counterfactual analysis
% 1. determine the output and prices and welfare under incomplete info. 
% 2. repeat #1 for complete info. 
% 3. compare #1 and #2. 
% 4. Repeat 1-3 several times. 
% 5. Average #4. 

% date created/modified June 27th 2022. 
clear
clc

thisFile = mfilename('fullpath');
thisDir = fileparts(thisFile);
cd(fileparts(fileparts(thisDir)));

load("output/estimates/ESTIMATES.mat") % file name is theta_mle

rng(1234)

%% Define the Model parameters using theta_mle 

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 %U~f_U =Truncated Normal(mu_u, sigma_u^2) truncated at u_low=theta_mle(20)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
ng =6; % number of groups
f_u = makedist('Normal','mu',theta_mle(3),'sigma',sqrt(theta_mle(4)));
f_u = truncate(f_u, theta_mle(end),Inf);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%f_{W|U}(.| U=u)
% W|U=u ~ w_bar x (2x Beta(a_w(u), a_w(u))-1) where 
% a_w(u) = exp(a_1_tilde + a_2_tilde x u) 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

a_w_u = @(u) exp(theta_mle(2*ng+5)+theta_mle(2*ng+6)*u);

% define Beta 
f_w_u = @(u)makedist('Beta', a_w_u(u),a_w_u(u));

% draw w's for all three values of u
w_draws_given_u = @(u) theta_mle(19)*(2.*random(f_w_u(u),1,1)-1);



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% V~f_{V}= w_bar x (Beta(a_i, b_i)+1)
% where i=1,...6 groups. 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% define the density parameters
a= theta_mle(5:ng+4); % Beta parameter a
b = theta_mle(ng+5:2*ng+4);% Beta parameter b
f_v_temp = @(grp_id) makedist("Beta", a(grp_id), b(grp_id)); %grp_id \in ng is group id
v_draw = @(grp_id, n_draw) theta_mle(19).*(random(f_v_temp(grp_id),n_draw,1)+1);
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%demand
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
I =20; % number of firms
Groups = [5, 4, 3, 5, 2, 1]; % firms within a group are symmetric. 


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% asymmetric information output
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% define the mean and variance of costs (also used in Figure 5). 
mu_vi       = theta_mle(19)*(a./(a+b)+1)';
variance_vi = theta_mle(19)^2*(a.*b./((a+b).^2.*(a+b+1)))';

% define sum of mean costs of everyone except i \sum_{j\neq i} \mu_{v_j}
number_of_firms_group = [5,4,3,5,2,1]; % how many firms in each group
 mu_vi_temp = [];
 for jj=1:ng
     temp = repmat(mu_vi(jj), number_of_firms_group(jj),1); 
     mu_vi_temp=[mu_vi_temp;temp];
 end

 Sum_mu_v_not_i = zeros(I,1);
for jj=1:I
    temp = mu_vi_temp;
    temp(jj)=[]; % remove the jjth row
    Sum_mu_v_not_i(jj) = sum(temp);
end

beta   = theta_mle(1);
lambda = theta_mle(2); 
 
q_i = @(grp_id, k,v_i, u, w)...
    1/(lambda+(I+1)*beta)*...
    (...
    u-w-1/(lambda+beta)*...
    ((lambda+I*beta)*mu_vi(grp_id)-...
    beta*Sum_mu_v_not_i(k))...
    )...
    -(v_i-mu_vi(grp_id))/(lambda+2*beta);
  

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% complete information output use complete_informaton.m
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
n_draws=1000;
Q = cell(n_draws,1); %save quantities
U_used = zeros(n_draws,1);  %save demand intercept
W_used = zeros(n_draws,1);  % common cost shock 
V_used = cell(n_draws,1);   % individual cost shock
Price = zeros(n_draws,2); % save eqm prices
Profit = cell(n_draws,1);
n=1;
e=1;
%check =0; % checks assumption 2-2
addpath('code/4_counterfactual');
while n<=n_draws % sometimes draws can get negative qt. discard.
        temp =zeros(I,2); %save output col 1/2 are incomplete/complete info
        U = random(f_u,1);% draw u
        W = w_draws_given_u(U); % draw w given u
        V =[]; %save mc
        for gg=1:ng
            V = [V;v_draw(gg,Groups(gg))]; % draw v
        end
        for kk=1:I
            if kk<=5
                temp(kk,1)=q_i(1,kk,V(kk), U, W);
            elseif kk>5 && kk<= 9
                temp(kk,1)=q_i(2,kk,V(kk), U, W);
            elseif kk>9 && kk<= 12
                temp(kk,1)=q_i(3,kk,V(kk), U, W);
            elseif kk>12 && kk<=17
                temp(kk,1)=q_i(4,kk,V(kk), U, W);
            elseif kk>17 && kk<=19
                temp(kk,1)=q_i(5,kk,V(kk), U, W);
            else
                temp(kk,1)=q_i(6,kk,V(kk), U, W);
            end
        end    

    temp(:,2) = complete_information(V, U, W, beta, lambda, I);

    temp1 = U-beta*sum(temp(:,1)); % price under incomplete info
    temp2= U-beta*sum(temp(:,2)); % price under complete info
    temp3 = (temp1 -(V+W)).*temp(:,1); %profit incomplete info
    temp4 = (temp2 -(V+W)).*temp(:,2); % profit complete info


    if sum(temp(:,1)>0)==I && temp1>0 && temp2>0 && sum(temp3>0)==I && sum(temp4>0)==I
        U_used(n,1) = U;
        W_used(n,1) = W; 
        V_used{n,1} = V; 
        Q{n,1} = temp;
        Price(n,:)=[temp1, temp2];
        Profit{n,1} = [temp3,temp4];
        n = n+1;
    end
    e=e+1
end
%%
 DWL   = zeros(n_draws,1); % calculate deadweight loss from moving from complete to incomplete
 CS = zeros(n_draws,2); % save consumer surplus
 for n=1:n_draws

    %%%% Incomplete Information 
    % reduction in cs (using formula for trapezoid) (a+b)/2*h
    cs_reduction = (sum(Q{n,1}(:,1)) ... %total output at price
        + (-(min(V_used{n,1})+W_used(n)-U_used(n))/beta))/2 ...% output at minimum cost 
    .*(Price(n,1)-min(V_used{n,1})-W_used(n)); % height price minus min cost
    
    % producer surplus
    PS = sum(Profit{n,1}(:,1)); 
    DWL(n,1) = cs_reduction-PS; 
     
    %%%% Complete Information 
    cs_reduction = (sum(Q{n,1}(:,2)) ... %total output at price
        + (-(min(V_used{n,1})+W_used(n)-U_used(n))/beta))/2 ...% output at minimum cost 
    .*(Price(n,2)-min(V_used{n,1})-W_used(n)); % height price minus min cost
    
    % producer surplus
    PS = sum(Profit{n,1}(:,2)); 

    DWL(n,2) = cs_reduction-PS; 
 end

 %% Deadweight Loss
 dwl=(DWL(:,1)-DWL(:,2))./DWL(:,1); % move from incomplete to complete lowers dwl
 dwl_mean=mean(dwl);
 
 %% Table to compare factual and counterfactual 
 Q_draws=zeros(I,n_draws,2); %firms x ndraws x information
 Q_average = zeros(I,2);
 Q_median  = zeros(I,2); 
 Q_min     = zeros(I,2); 
 Q_std     = zeros(I,2); 
 for ii=1:I
     temp =[];
     for kk=1:n_draws
         temp_ = Q{kk,1}(ii,:);
         temp = [temp; temp_]; 
     end
     Q_draws(ii,:,:)=temp;
     Q_average(ii,:)= mean(temp);
     Q_median(ii,:)= median(temp);
     Q_min(ii,:)= min(temp);
     Q_std(ii,:)= var(temp);
     clear temp
 end
 Q_average = mean(Q_average); %cols are (incomplete, complete info)
 Q_median = mean(Q_median);
 Q_std     = mean(Q_std);

 Price_average = mean(Price); %cols are (incomplete, complete info)
 change_price = mean(100*(Price(:,1)-Price(:,2))./Price(:,1));

 % save the final results. 
save('output/estimates/CF_results.mat', 'Q_average', 'Price_average', 'change_price');
 